// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package v1

import (
	"fmt"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// NPUClusterPolicySpec defines the desired state of NPUClusterPolicy
type NPUClusterPolicySpec struct {
	// Operator configuration
	Operator OperatorSpec `json:"operator"`
	// Daemonset configuration
	Daemonsets DaemonsetsSpec `json:"daemonsets"`
	// Ascend Driver componet spec
	Driver DriverSpec `json:"driver"`
	// Ascend Docker Runtime component spec
	OCIRuntime OCIRuntimeSpec `json:"ociRuntime"`
	// Volcano component spec
	VCScheduler  SchedulerSpec  `json:"vcscheduler"`
	VCController ControllerSpec `json:"vccontroller"`
	// Ascend Operator component spec
	Trainer TrainerSpec `json:"trainer"`
	// Ascend Device Plugin component spec
	DevicePlugin DevicePluginSpec `json:"devicePlugin"`
	// Ascend NodeD component spec
	NodeD NodeDSpec `json:"nodeD"`
	// Ascend NPU Exporter component spec
	Exporter ExporterSpec `json:"exporter"`
	// Ascend ClusterD component spec
	ClusterD ClusterDSpec `json:"clusterd"`
	// Ascend Resilience Controller component spec
	RSController RSControllerSpec `json:"rscontroller"`
	// MindIO component spec
	MindIOTFT MindIOTFTSpec `json:"mindiotft"`
	// MindIO component spec
	MindIOACP MindIOACPSpec `json:"mindioacp"`
}

// CommandSpec defines the properties for container
type CommandSpec struct {
	Command []string `json:"command,omitempty"`
	Args    []string `json:"args,omitempty"`
}

// ImageSpec defines the properties for an image
type ImageSpec struct {
	// Image registry
	Registry string `json:"registry,omitempty"`

	// Image repository
	Repository string `json:"repository,omitempty"`

	// Image tag
	Tag string `json:"tag,omitempty"`

	// Image pull policy
	// +default="IfNotPresent"
	ImagePullPolicy string `json:"imagePullPolicy,omitempty"`

	// Image pull secrets
	// +default=[]
	ImagePullSecrets []string `json:"imagePullSecrets,omitempty"`
}

// Name returns the complete image name
func (s ImageSpec) Name() string {
	registry := s.Registry
	repo := s.Repository
	tag := s.Tag
	var image string
	if len(repo) == 0 {
		return ""
	}
	if len(registry) == 0 && len(tag) == 0 {
		if len(repo) != 0 {
			image = repo
		}
	} else {
		if len(registry) == 0 {
			image = fmt.Sprintf("%s:%s", repo, tag)
		} else {
			image = fmt.Sprintf("%s/%s:%s", registry, repo, tag)
		}
	}
	return image
}

// OperatorSpec defines the configuration for the operator
type OperatorSpec struct {
	// Name of CRI Runtime Handler
	// +kubebuilder:default=ascend
	RuntimeClass string `json:"runtimeClass,omitempty"`

	// Image configuration
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
}

// DaemonsetsSpec defines common configuration for all Daemonsets
type DaemonsetsSpec struct {
	// Common labels
	// +optional
	Labels map[string]string `json:"labels,omitempty"`
	// Common annotations
	// +optional
	Annotations map[string]string `json:"annotations,omitempty"`
	// Common tolerations
	// +optional
	// +listType=atomic
	Tolerations []corev1.Toleration `json:"tolerations,omitempty"`

	// Common image
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
}

// OCIRuntimeSpec defines the properties for Ascend Docker Runtime deployment
type OCIRuntimeSpec struct {
	// Whether to manage Docker Runtime with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Interval for health checking
	// +default=300
	// +kubebuilder:validation:Minimum=60
	// +kubebuilder:validation:Maximum=1800
	Interval int32 `json:"interval,omitempty"`
	// Image configuration
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
	// Init runtime image configuration
	// +default={}
	InitRuntimeImage ImageSpec `json:"initRuntimeImageSpec,omitempty"`
	// Init config image configuration
	// +default={}
	InitConfigImage ImageSpec `json:"initConfigImageSpec,omitempty"`
}

// SchedulerSpec defines the properties for Volcano Scheduler deployment
type SchedulerSpec struct {
	// Whether to manage Volcano with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Resources required by scheduler container
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image              ImageSpec   `json:"imageSpec,omitempty"`
	Command            CommandSpec `json:"commandSpec,omitempty"`
	SchedulerConfigMap string      `json:"schedulerConfigMap,omitempty"`
}

// ControllerSpec defines the properties for Volcano Controller deployment
type ControllerSpec struct {
	// Whether to manage Volcano with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Resources required by controller container
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image   ImageSpec   `json:"imageSpec,omitempty"`
	Command CommandSpec `json:"commandSpec,omitempty"`
}

// MindIOTFTSpec defines the properties for MindIO Training Fault Tolerance daemonset
type MindIOTFTSpec struct {
	// whether to manage MindIO-TFT with NPU Operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Image configuration
	Image ImageSpec `json:"imageSpec,omitempty"`
}

// MindIOACPSpec defines the properties for MindIO Async Checkpoint Persistence daemonset
type MindIOACPSpec struct {
	// whether to manage MindIO-ACP with NPU Operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Image configuration
	Version string    `json:"version,omitempty"`
	Image   ImageSpec `json:"imageSpec,omitempty"`
}

// ResourceRequirements defines the compute resource requirements
type ResourceRequirements struct {
	// Same as corev1.ResourceRequirements.Limits
	Limits corev1.ResourceList `json:"limits,omitempty"`
	// Same as corev1.ResourceRequirements.Requests
	Requests corev1.ResourceList `json:"requests,omitempty"`
}

// VolumesSpec defines the volume path of container
type VolumesSpec struct {
	Volumes      []corev1.Volume      `json:"volumes,omitempty"`
	VolumeMounts []corev1.VolumeMount `json:"volumeMounts,omitempty"`
}

// TrainerSpec defines the properties for Ascend Operator deployment
type TrainerSpec struct {
	// Whether to manage Ascend Operator with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/ascend-operator/ascend-operator.log"}
	LogRotate LogRotate            `json:"logRotate,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
	// Init image configuration
	// +default={}
	InitImage ImageSpec   `json:"initImageSpec,omitempty"`
	Command   CommandSpec `json:"commandSpec,omitempty"`
}

// LogRotate defines the configuration for log rotation
type LogRotate struct {
	// Log level
	// +default="info"
	LogLevel LogLevel `json:"logLevel,omitempty"` // NOTE: -1 ~ 3, debug ~ critical
	// Log file
	LogFile string `json:"logFile,omitempty"`
	// Number of rotations before log file is removed
	// +default=30
	// +kubebuilder:validation:Minimum=1
	// +kubebuilder:validation:Maximum=30
	Rotate int32 `json:"rotate,omitempty"` // NOTE: -maxBackups
	// Maximum days to keep before log file is removed
	// +default=7
	// +kubebuilder:validation:Minimum=7
	// +kubebuilder:validation:Maximum=700
	MaxAge int32 `json:"maxAge,omitempty"`
	// Enable/disable compression of rotated log files
	// +default=false
	Compress bool `json:"compress,omitempty"` // NOTE: -isCompress
}

// +enum
// +kubebuilder:validation:Enum=debug;info;warning;error;critical

// LogLevel defines the log level of NPU operator components
type LogLevel string

// Value returns the log level value
func (l LogLevel) Value() int {
	switch l {
	case LogDebug:
		return LogDebugValue
	case LogInfo:
		return LogInfoValue
	case LogWarning:
		return LogWarningValue
	case LogError:
		return LogErrorValue
	case LogCritical:
		return LogCriticalValue
	default:
		panic(fmt.Sprintf("unsupported LogLevel: %v", l))
	}
}

const (
	// LogDebug indicates debug log level
	LogDebug LogLevel = "debug"
	// LogInfo indicates info log level
	LogInfo LogLevel = "info"
	// LogWarning indicates warning log level
	LogWarning LogLevel = "warning"
	// LogError indicates error log level
	LogError LogLevel = "error"
	// LogCritical indicates critical log level
	LogCritical LogLevel = "critical"

	// LogDebugValue indicates value of debug log level
	LogDebugValue = -1
	// LogInfoValue indicates value of info log level
	LogInfoValue = 0
	// LogWarningValue indicates value of warning log level
	LogWarningValue = 1
	// LogErrorValue indicates value of error log level
	LogErrorValue = 2
	// LogCriticalValue indicates value of critical log level
	LogCriticalValue = 3
)

// DevicePluginSpec defines the properties for Ascend Device Plugin deployment
type DevicePluginSpec struct {
	// Whether to manage Device Plugin with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/devicePlugin/devicePlugin.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate `json:"logRotate,omitempty"`

	// Image configuration
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
	// Init image configuration
	// +default={}
	InitImage ImageSpec            `json:"initImageSpec,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Container configuration
	// +default={}
	Command    CommandSpec `json:"commandSpec,omitempty"`
	VolumeSpec VolumesSpec `json:"volumeSpec,omitempty"`
}

// NodeDSpec defines the properties for Ascend NodeD deployment
type NodeDSpec struct {
	// Whether to manage NodeD with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Interval for sending heartbeat messages
	// +default=5
	// +kubebuilder:validation:Minimum=1
	// +kubebuilder:validation:Maximum=300
	HeartbeatInterval int32 `json:"heartbeatInterval,omitempty"`
	// Interval for scanning hardware failure on nodes
	// +default=60
	// +kubebuilder:validation:Minimum=60
	// +kubebuilder:validation:Maximum=600
	PollInterval int32 `json:"pollInterval,omitempty"` // NOTE: monitorPeriod
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/noded/noded.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate            `json:"logRotate,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image   ImageSpec   `json:"imageSpec,omitempty"`
	Command CommandSpec `json:"commandSpec,omitempty"`
}

// ExporterSpec defines the properties for Ascend NPU Exporter deployment
type ExporterSpec struct {
	// Whether to manage NPU Exporter with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/npu-exporter/npu-exporter.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate            `json:"logRotate,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image   ImageSpec   `json:"imageSpec,omitempty"`
	Command CommandSpec `json:"commandSpec,omitempty"`
}

// DriverSpec defines the properties for Ascend NPU Driver daemonset
type DriverSpec struct {
	// Whether to manage NPU Driver with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/driver/driver.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate `json:"logRotate,omitempty"`
	// Version of hdk.run
	Version string `json:"version,omitempty"`
	// Image configuration
	// +default={}
	Image ImageSpec `json:"imageSpec,omitempty"`
	// Init image configuration
	// +default={}
	InitImage ImageSpec `json:"initImageSpec,omitempty"`
}

// ClusterDSpec defines the properties for ClusterD deployment
type ClusterDSpec struct {
	// Whether to manage ClusterD with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/clusterd/clusterd.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate            `json:"logRotate,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image   ImageSpec   `json:"imageSpec,omitempty"`
	Command CommandSpec `json:"commandSpec,omitempty"`
}

// RSControllerSpec defines the properties for Resilience Controller deployment
type RSControllerSpec struct {
	// Whether to manage Resilience Controller with NPU operator
	// +default=true
	Managed bool `json:"managed,omitempty"`
	// Log rotation
	// +default={"logFile":"/var/log/mindx-dl/resilience-controller/run.log"}
	// +kubebuilder:validation:XValidation:rule="!has(self.compress) || !self.compress",message="Compress is unsupported"
	LogRotate LogRotate            `json:"logRotate,omitempty"`
	Resources ResourceRequirements `json:"resources,omitempty"`
	// Image configuration
	// +default={}
	Image   ImageSpec   `json:"imageSpec,omitempty"`
	Command CommandSpec `json:"commandSpec,omitempty"`
}

// NPUClusterPolicyStatus defines the observed state of NPUClusterPolicy
type NPUClusterPolicyStatus struct {
	// State of NPUClusterPolicy
	Phase Phase `json:"phase"`
	// Namespace where the operator is installed
	Namespace string `json:"namespace,omitempty"`
	// Observations of NPUClusterPolicy's current state.
	// Known .status.conditions.type are: "Ready" and "Error"
	// +patchMergeKey=type
	// +patchStrategy=merge
	// +listType=map
	// +listMapKey=type
	Conditions []metav1.Condition `json:"conditions,omitempty" patchStrategy:"merge" patchMergeKey:"type"`
	// Statuses for all NPU operator components, one entry per component
	// +listType=atomic
	ComponentStatuses []ComponentStatus `json:"componentStatuses,omitempty"`
}

// UpdateConditions defines the function that update conditions of NPUClusterPolicyStatus.
func (s *NPUClusterPolicyStatus) UpdateConditions(oldType string, newType string, reason string, message string) bool {
	changed := meta.SetStatusCondition(&s.Conditions, metav1.Condition{
		Type:    newType,
		Status:  metav1.ConditionTrue,
		Reason:  reason,
		Message: message,
	})
	changed = meta.SetStatusCondition(&s.Conditions, metav1.Condition{
		Type:   oldType,
		Status: metav1.ConditionFalse,
		Reason: newType,
	}) || changed
	return changed
}

// UpdateComponentStatus add or update component status of NPUClusterPolicyStatus.
func (s *NPUClusterPolicyStatus) UpdateComponentStatus(name string, state ComponentState) bool {
	if cs := s.GetComponentStatus(name); cs != nil {
		if cs.State == state {
			return false
		}
		cs.PrevState = cs.State
		cs.State = state
	} else {
		s.ComponentStatuses = append(s.ComponentStatuses, ComponentStatus{Name: name, State: state})
	}
	return true
}

// GetComponentStatus returns the component status of NPUClusterPolicyStatus by component name.
func (s *NPUClusterPolicyStatus) GetComponentStatus(name string) *ComponentStatus {
	for i := range s.ComponentStatuses {
		cs := &s.ComponentStatuses[i]
		if cs.Name == name {
			return cs
		}
	}
	return nil
}

// +enum
// +kubebuilder:validation:Enum=Ready;NotReady

// Phase defines the current state of NPUClusterPolicy
type Phase string

const (
	// PolicyReady indicates all components of NPUClusterPolicy are ready
	PolicyReady Phase = "Ready"
	// PolicyNotReady indicates some/all components of NPUClusterPolicy are not ready
	PolicyNotReady Phase = "NotReady"
)

// ComponentStatus defines the current status for an NPU operator component
type ComponentStatus struct {
	// Name of the component
	Name string `json:"name"`
	// State of the component
	State ComponentState `json:"state,omitempty"`
	// Previous state of the component for debugging purpose
	PrevState ComponentState `json:"prevState,omitempty"`
}

// ComponentState defines the state of an NPU operator component
type ComponentState struct {
	// Phase of the component state
	// +default="pending"
	Phase ComponentPhase `json:"type,omitempty"`
	// Unique, one-word, CamelCase reason for the component state
	Reason string `json:"reason,omitempty"`
	// Human-readable message indicating details about the component state
	Message string `json:"message,omitempty"`
}

// +enum
// +kubebuilder:validation:Enum=pending;deploying;running;terminated;unmanaged

// ComponentPhase defines the current state of an NPU operator component
type ComponentPhase string

const (
	// ComponentPending indicates the component is pending deployment, or is failed to deploy and pending redeployment
	ComponentPending ComponentPhase = "pending"
	// ComponentDeploying indicates the component is being deployed
	ComponentDeploying ComponentPhase = "deploying"
	// ComponentRunning indicates the component is successfully deployed and running
	ComponentRunning ComponentPhase = "running"
	// ComponentTerminated indicates the component is terminated and pending restart
	ComponentTerminated ComponentPhase = "terminated"
	// ComponentUnmanaged indicates the component is not managed by the operator
	ComponentUnmanaged ComponentPhase = "unmanaged"
)

// +genclient
// +genclient:nonNamespaced
// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:resource:scope=Cluster
// +kubebuilder:printcolumn:name="Status",type=string,JSONPath=`.status.phase`,priority=0
// +kubebuilder:printcolumn:name="Age",type=string,JSONPath=`.metadata.creationTimestamp`,priority=0
// +kubebuilder:validation:XValidation:rule="self.metadata.name == 'cluster'",message="Only one NPUClusterPolicy with name 'cluster' can be created in a cluster"

// NPUClusterPolicy is the Schema for the npuclusterpolicies API
type NPUClusterPolicy struct {
	metav1.TypeMeta   `json:",inline"`
	metav1.ObjectMeta `json:"metadata,omitempty"`

	Spec   NPUClusterPolicySpec   `json:"spec,omitempty"`
	Status NPUClusterPolicyStatus `json:"status,omitempty"`
}

// +kubebuilder:object:root=true

// NPUClusterPolicyList contains a list of NPUClusterPolicy
type NPUClusterPolicyList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata,omitempty"`
	Items           []NPUClusterPolicy `json:"items"`
}

func init() {
	SchemeBuilder.Register(&NPUClusterPolicy{}, &NPUClusterPolicyList{})
}
